# -*- coding: utf-8 -*-
"""
ZMQ管理器
负责ZMQ上下文和Socket的创建、管理和销毁
"""

from sre_parse import SUCCESS
from typing import Any, Dict, Tuple
import errno
import json
import zmq
import zmq.asyncio
from common.utils import LogUtil
from common.config import ServiceConfig

logger = LogUtil.get_logger(__name__)

class ZmqServer:
    """ZMQ管理器类"""
    
    def __init__(self, host: str):
        """初始化ZMQ管理器
        
        Args:
            host (str): 绑定的主机地址
        """
        self._host = host
        self._zmq_context = None
        self._aio_zmq_context = None
        self._rep_socket = None
        self._pub_socket = None
        self._aio_socket = None
        
        self._zmq_context = zmq.Context()
        self._reset_rep_socket()
        self._reset_pub_socket()

        self._aio_zmq_context = zmq.asyncio.Context()
        self._reset_aio_socket()

        logger.info("ZMQ管理器初始化完成")
    
    def _reset_rep_socket(self):
        """重置ZMQ REP Socket"""
        if self._rep_socket:
            self._rep_socket.close()
            self._rep_socket = None
        self._rep_socket = self._zmq_context.socket(zmq.REP)
        self._rep_socket.bind(f"tcp://{self._host}:{ServiceConfig.ZMQ_REP_PORT}")
        self._rep_socket.setsockopt(zmq.RCVTIMEO, ServiceConfig.ZMQ_TIMEOUT)
        self._rep_socket.setsockopt(zmq.SNDTIMEO, ServiceConfig.ZMQ_TIMEOUT)
        logger.info("ZMQ REP Socket已重置")
    
    def _reset_pub_socket(self):
        """重置ZMQ PUB Socket"""
        if self._pub_socket:
            self._pub_socket.close()
            self._pub_socket = None
        self._pub_socket = self._zmq_context.socket(zmq.PUB)
        self._pub_socket.bind(f"tcp://{self._host}:{ServiceConfig.ZMQ_PUB_PORT}")
        self._pub_socket.setsockopt(zmq.RCVTIMEO, ServiceConfig.ZMQ_TIMEOUT)
        self._pub_socket.setsockopt(zmq.SNDTIMEO, ServiceConfig.ZMQ_TIMEOUT)
        logger.info("ZMQ PUB Socket已重置")
    
    def _reset_aio_socket(self):
        """重置ZMQ AIO Socket"""
        if self._aio_socket:
            self._aio_socket.close()
            self._aio_socket = None
        self._aio_socket = self._aio_zmq_context.socket(zmq.REP)
        self._aio_socket.bind(f"tcp://{self._host}:{ServiceConfig.ZMQ_AIO_PORT}")
        self._aio_socket.setsockopt(zmq.RCVTIMEO, ServiceConfig.ZMQ_TIMEOUT)
        self._aio_socket.setsockopt(zmq.SNDTIMEO, ServiceConfig.ZMQ_TIMEOUT)
        logger.info("ZMQ AIO Socket已重置")

    def destroy(self):
        """销毁ZMQ资源"""
        logger.info("销毁ZMQ管理器")
        # 关闭Socket
        if self._rep_socket:
            self._rep_socket.close()
            self._rep_socket = None
        if self._pub_socket:
            self._pub_socket.close()
            self._pub_socket = None
        if self._aio_socket:
            self._aio_socket.close()
            self._aio_socket = None
        # 终止异步上下文
        if self._aio_zmq_context:
            self._aio_zmq_context.term()
            self._aio_zmq_context = None
        # 终止同步上下文
        if self._zmq_context:
            self._zmq_context.term()
            self._zmq_context = None
        logger.info("ZMQ管理器已销毁")
    
    def __del__(self):
        """析构函数"""
        self.destroy()

    def recv_json(self) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        request = None
        success = False
        try:
            request = self._rep_socket.recv_json()
            success = True
            logger.debug(f"接收UiTest请求成功")
        except zmq.error.Again as e:
            # logger.debug(f"接收UiTest请求超时: {e}")
            pass
        except zmq.error.ZMQError as e:
            if e.errno == errno.EINTR:
                logger.debug(f"接收UiTest请求被信号中断: {e}")
                pass
            else:
                logger.error(f"接收UiTest请求发生ZMQ错误: {e}")
                self._reset_rep_socket()
        except json.JSONDecodeError as e:
            logger.error(f"解析UiTest请求失败: {e}")
            success = True
        except Exception as e:
            logger.error(f"接收UiTest请求失败: {e}")
            self._reset_rep_socket()
        return (request, success)

    async def async_recv_json(self):
        request = None
        success = False
        try:
            request = await self._aio_socket.recv_json()
            success = True
            logger.info("接收信令消息成功")
        except zmq.error.Again:
            # logger.info(f"接收信令消息超时")
            pass
        except zmq.error.ZMQError as e:
            if e.errno == errno.EINTR:
                logger.info(f"接收信令消息被信号中断: {e}")
                pass
            else:
                logger.error(f"接收信令消息发生ZMQ错误: {e}")
                self._reset_aio_socket()
        except json.JSONDecodeError as e:
            logger.error(f"解析信令消息失败: {e}")
            success = True
        except Exception as e:
            logger.error(f"接收信令消息失败: {e}")
            self._reset_aio_socket()
        return (request, success)

    def send_json(self, response: Dict[str, Any]) -> bool:
        """发送UiTest响应"""
        while True:
            try:
                self._rep_socket.send_json(response)
                logger.debug("发送UiTest响应成功")
                return True
            except zmq.error.Again as e:
                # logger.debug(f"发送UiTest响应超时: {e}")
                continue
            except zmq.error.ZMQError as e:
                if e.errno == errno.EINTR:
                    logger.debug(f"发送UiTest响应被信号中断: {e}")
                    continue
                else:
                    logger.error(f"发送UiTest响应发生ZMQ错误: {e}")
                    self._reset_rep_socket()
                    return False
            except Exception as e:
                logger.error(f"发送UiTest响应失败: {e}")
                self._reset_rep_socket()
                return False
    
    async def async_send_json(self, response: Dict[str, Any]) -> bool:
        while True:
            try:
                await self._aio_socket.send_json(response)
                logger.info("发送信令消息成功")
                return True
            except zmq.error.Again:
                # 发送超时
                logger.info("发送信令消息超时")
                continue
            except zmq.error.ZMQError as e:
                if e.errno == errno.EINTR:
                    logger.info(f"发送信令消息被信号中断: {e}")
                    continue
                else:
                    logger.error(f"发送信令消息发生ZMQ错误: {e}")
                    self._reset_aio_socket()
                    return False
            except Exception as e:
                logger.error(f"发送信令消息失败: {e}")
                self._reset_aio_socket()
                return False

    def send_multipart(self, msg_parts) -> bool:
        try:
            self._pub_socket.send_multipart(msg_parts)
            logger.info(f"发送设备事件成功")
            return True
        except Exception as e:
            self._reset_pub_socket()
            logger.error(f"发送设备事件失败: {e}")
            return False
